from flask import Flask, request, jsonify
import joblib
import pandas as pd
import pickle
import traceback
import torch
import os

os_path = os.getcwd()
print (os_path)
app = Flask(__name__)
cur_path = '/home/szr/.A003_AI_Clock_Server/aiclock.server/'

from SimpleNN import SimpleNN

# 设置计算设备为 GPU
is_cuda = torch.cuda.is_available()
device = torch.device("cuda" if is_cuda else "cpu")

# 加载模型
try:
    model = joblib.load(cur_path + 'model_GRU1.pkl')
    model.to(device)  # 将模型移动到 GPU 设备
    model.eval()  # 设置模型为评估模式
    print("Model loaded successfully")
except (OSError, joblib.JoblibException) as e:
    print(f"Error loading model: {e}")
    model = None
    print(traceback.print_exc())

# 加载标准化器和特征列
try:
    with open(cur_path + 'scaler.pkl', 'rb') as f:
        scaler = pickle.load(f)
    with open(cur_path + 'model_columns.pkl', 'rb') as f:
        model_columns = pickle.load(f)
    print("Scaler and model columns loaded successfully")
except (OSError, pickle.PickleError) as e:
    print(f"Error loading scaler or model columns: {e}")
    scaler = None
    model_columns = None

@app.route('/')
def home():
    return "Welcome to the model API!"

@app.route('/predict', methods=['POST', 'OPTIONS'])
def predict():
    print("get")

    if model is None or scaler is None or model_columns is None:
        return jsonify({'error': 'Model, scaler, or model columns are not available'}), 500

    if request.method == 'OPTIONS':
        print('Received OPTIONS request')
        response = app.make_default_options_response()
        headers = response.headers

        headers['Access-Control-Allow-Origin'] = '*'
        headers['Access-Control-Allow-Methods'] = 'POST'
        headers['Access-Control-Allow-Headers'] = 'Content-Type'

        return response
    
    if request.method == 'POST':
        print('Received POST request')

        try:
            data = request.get_json()  # 获取POST请求中的JSON数据
            features = data['features']

            # 创建一个 DataFrame 与训练时的特征相匹配
            df = pd.DataFrame([features])

            # 进行独热编码
            df_encoded = pd.get_dummies(df, columns=['起点', '目的地', '当前时间', '路径'])

            # 确保数据框的列顺序与训练时的顺序一致
            df_encoded = df_encoded.reindex(columns=model_columns, fill_value=0)

            # 打印独热编码后的 DataFrame 形状
            print(f"Encoded DataFrame shape: {df_encoded.shape}")

            # 进行标准化
            X_scaled = scaler.transform(df_encoded)

            # 打印标准化后的数据形状
            print(f"Scaled Data shape: {X_scaled.shape}")

            # 进行预测
            inputs = torch.from_numpy(X_scaled).float().to(device)  # 确保输入张量在同一个设备上
            with torch.no_grad():
                prediction = model(inputs).cpu().numpy()  # 将预测结果移回CPU以便处理
            print(prediction)
            response = jsonify({'prediction': prediction.tolist()})

            response.headers['Access-Control-Allow-Origin'] = '*'
            response.headers['Access-Control-Allow-Methods'] = 'POST'
            response.headers['Access-Control-Allow-Headers'] = 'Content-Type'
            
            return response
            
        except Exception as e:
            print('Error:', e)  # 添加日志
            print(traceback.print_exc())  # 打印详细的错误信息
            response = jsonify({'error': str(e)})
            response.headers['Access-Control-Allow-Origin'] = '*'
            response.headers['Access-Control-Allow-Methods'] = 'POST'
            response.headers['Access-Control-Allow-Headers'] = 'Content-Type'
            return response, 400  # 返回错误信息

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=7861)
